﻿#pragma once

#include <coroutine>
#include <future>

namespace xi
{

namespace utils
{

template <typename Ty> struct Generator
{
    // 协程控制结构
    struct promise_type
    {
        auto get_return_object()
        {
            return Generator{*this};
        }

        // 创建协程同时暂停
        auto initial_suspend() noexcept
        {
            // suspend_always 表示暂停
            // suspend_never 表示不暂停
            return std::suspend_always{};
        }

        // 协程完成后暂停
        auto final_suspend() noexcept
        {
            return std::suspend_always{};
        }

        // co_return 关键字的返回值（void 类型）
        void return_void() noexcept
        {
        }

        // co_yield 关键字的返回值，返回后暂停
        auto yield_value(const Ty &value) noexcept
        {
            // addressof 取地址，防止 & 运算符被重载
            current_value = std::addressof(value);
            return std::suspend_always{};
        }

        // 当协程异常时，调用此函数
        void unhandled_exception() noexcept
        {
            exception = std::current_exception();
        }

        const Ty *current_value;
        std::exception_ptr exception;
    };

    // 协程句柄，管理协程生命周期
    using handle = std::coroutine_handle<promise_type>;

  private:
    handle m_handle{nullptr};

  public:
    explicit Generator(promise_type &prom) noexcept : m_handle(handle::from_promise(prom))
    {
    }

    // 移动构造，交换句柄
    Generator(Generator &&rhs) noexcept : m_handle(std::exchange(rhs.m_handle, nullptr))
    {
    }

    // 移动赋值，交换句柄
    Generator &operator=(Generator &&right) noexcept
    {
        m_handle = std::exchange(right.m_handle, nullptr);
        return *this;
    }

    ~Generator()
    {
        if (m_handle)
            m_handle.destroy();
    }

    class iterator
    {
      private:
        handle m_coro{nullptr};

      public:
        using iterator_category = std::input_iterator_tag; // 标记输入迭代器
        using difference_type = std::ptrdiff_t;            // 计算迭代器距离
        using value_type = Ty;
        using reference = const Ty &;
        using pointer = const Ty *;

        // 用于构建空迭代器
        iterator() = default;

        explicit iterator(handle coro) noexcept : m_coro(coro)
        {
        }

        // 每次迭代都恢复协程，直到协程完成
        iterator &operator++()
        {
            m_coro.resume();
            if (m_coro.done())
                m_coro = nullptr;
            return *this;
        }

        iterator &operator++(int)
        {
            return (*this)++;
        }

        bool operator==(const iterator &right) const noexcept
        {
            return m_coro == right.m_coro;
        }

        bool operator!=(const iterator &right) const noexcept
        {
            return !(*this == right);
        }

        reference operator*() const noexcept
        {
            return *m_coro.promise().current_value;
        }

        pointer operator->() const noexcept
        {
            return m_coro.promise().current_value;
        }
    };

    iterator begin()
    {
        if (m_handle)
        {
            m_handle.resume();
            if (m_handle.done())
                return {};
        }

        return iterator{m_handle};
    }

    iterator end() noexcept
    {
        return {};
    }
};

} // namespace utils

} // namespace xi
